from skyfield.api import load, wgs84, EarthSatellite
from pytz import timezone
from datetime import datetime

ts = load.timescale()

shanghai = timezone("Asia/Shanghai")


class SatelliteTrack:
    def __init__(self):
        self.satellite_map = {}
        self.ts = ts

    def add_satellite_from_line_data(self, line1, line2, name):
        """
        从两行数据添加卫星
        :param line1:
        :param line2:
        :param name:
        :return:
        """
        satellite = EarthSatellite(line1=line1, line2=line2, name=name)
        self.satellite_map[name] = satellite
        return satellite

    def add_satellite_from_url(self, url="http://celestrak.org/NORAD/elements/stations.txt"):
        """
        从网络添加中添加卫星数据，新添加的数据会覆盖之前的老数据
        :param url:
        :return:
        """
        satellites = load.tle_file(url, reload=False)
        satellite_map = {sat.name: sat for sat in satellites}
        self.satellite_map.update(satellite_map)
        return satellite_map

    def transit(self, satellite_name, start_time, end_time, lat, lon, elevation_m=400, altitude_degrees=10):
        """
        # 计算卫星合适升起落下
        :param elevation_m:     海拔
        :param satellite_name:  卫星名字
        :param start_time:
        :param end_time:
        :param lat:
        :param lon:
        :param altitude_degrees:
        :return:
        """
        start_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(start_time))))
        end_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(end_time))))

        satellite = self.satellite_map[satellite_name]
        bluffton = wgs84.latlon(lat, lon, elevation_m)
        t, events = satellite.find_events(
            bluffton,
            start_time,
            end_time,
            altitude_degrees
        )
        de421 = load('de421.bsp')
        sunlit = satellite.at(t).is_sunlit(de421)  # 是否可观测
        difference = satellite - bluffton

        rise_down_list = []
        for ti, event, sunlit_flag in zip(t, events, sunlit):
            ti = ts.tt_jd(ti.tt)
            alt, az, distance = difference.at(ti).altaz()

            ti_str = ti.utc_iso()
            rise_down_list.append({
                # "time": ti_str,
                "time": int(datetime.fromisoformat(ti_str).timestamp() * 1000),
                "event": int(event),
                "degree_alt": int(alt.degrees),
                "degree_azimuth": int(az.degrees),
                "distance": distance.km,
                "sunlit_flag": bool(sunlit_flag)
            })
        return rise_down_list

    def get_position(self, satellite_name, start_time, end_time, lat, lon, elevation_m=100):
        """
        # 计算某一段时间内的观测者和卫星的观测角度
        :param elevation_m:海拔高度
        :param satellite_name:
        :param start_time:
        :param end_time:
        :param lat:
        :param lon:
        :return:
        """
        satellite = self.satellite_map[satellite_name]
        bluffton = wgs84.latlon(lat, lon, elevation_m)
        difference = satellite - bluffton
        event_list = []

        start_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(start_time))))
        end_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(end_time))))

        time_len = (end_time - start_time) * 3600 * 24
        for i in range(int(time_len)):
            ti = start_time + (i / (3600 * 24))
            pos = difference.at(ti)
            alt, az, distance = pos.altaz()

            # 计算多普勒
            # _, _, the_range, _, _, range_rate = pos.frame_latlon_and_rates(bluffton)

            ti_str = ti.utc_iso()

            event_list.append({
                "time": ti_str,
                "degree_alt": alt.degrees,
                "degree_azimuth": az.degrees,
                "distance": distance.km
            })
        return event_list

    def get_satellite_gps(self, satellite_name, start_time, end_time, step_time_s=1):
        """
        # 计算某一时刻卫星的GPS位置 星下点
        :param end_time:
        :param start_time:
        :param satellite_name:
        :param step_time_s:
        :return:
        """
        satellite = self.satellite_map[satellite_name]
        start_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(start_time))))
        end_time = ts.from_datetime(shanghai.localize(datetime.fromtimestamp(int(end_time))))

        gps_list = []
        time = start_time
        while end_time - time >= 0:
            # 地心坐标系的位置
            geocentric = satellite.at(time)
            lat, lon = wgs84.latlon_of(geocentric)
            dt, leap_second = time.utc_datetime_and_leap_second()
            gps_list.append({
                "time": dt,
                "lat": lat.degrees,
                "lon": lon.degrees
            })
            time = time + (step_time_s / (3600 * 24))
        return gps_list


if __name__ == "__main__":
    # https://celestrak.org/NORAD/elements/gp.php?GROUP=stations&FORMAT=tle
    # 空间站TLE
    line1 = "1 25544U 98067A   23326.88766336  .00027400  00000+0  49383-3 0  9995"
    line2 = "2 25544  51.6430 262.4949 0000489 314.5551 149.9978 15.49569765426402"
    satellite_name = "TianHe"
    satellite_track = SatelliteTrack()
    satellite_track.add_satellite_from_line_data(line1, line2, satellite_name)
    now = int(datetime.now().timestamp())
    res = satellite_track.transit(satellite_name, now, now + 24 * 3600, 30.5818, 104.05555, 10)
    print(res)
